package com.jxpanda.spring.module.auth.endpoint;

import com.jxpanda.infrastructure.core.tollkit.StringKit;
import com.jxpanda.spring.module.auth.config.properties.TokenMode;
import com.jxpanda.spring.module.auth.config.properties.TokenProperties;
import com.jxpanda.spring.module.auth.core.authentication.strategy.CollaborativeReactiveAuthenticationOAuth2StrategyResolver;
import com.jxpanda.spring.module.auth.core.authentication.token.CollaborativeAuthenticationToken;
import com.jxpanda.spring.module.auth.core.authentication.token.UserDetailsAuthenticationToken;
import com.jxpanda.spring.module.auth.core.token.ReactiveOAuth2TokenService;
import com.jxpanda.spring.module.auth.endpoint.handler.ReactiveAuthenticationFailureHandler;
import com.jxpanda.spring.module.auth.endpoint.handler.ReactiveAuthenticationSuccessHandler;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.lang.NonNull;
import org.springframework.security.core.context.ReactiveSecurityContextHolder;
import org.springframework.security.core.context.SecurityContextImpl;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.security.oauth2.client.registration.ReactiveClientRegistrationRepository;
import org.springframework.security.web.server.context.NoOpServerSecurityContextRepository;
import org.springframework.security.web.server.context.ServerSecurityContextRepository;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

@Slf4j
@RestController
@RequestMapping
@RequiredArgsConstructor
public class AuthEndpoint {

    public static final String REGISTRATION_ID = "registrationId";

    public static final String DEFAULT_PATH = "/auth/token/{" + REGISTRATION_ID + "}";

    private final ReactiveClientRegistrationRepository reactiveClientRegistrationRepository;

    private final CollaborativeReactiveAuthenticationOAuth2StrategyResolver collaborativeReactiveAuthenticationStrategyResolver;

    private final ReactiveOAuth2TokenService reactiveOAuth2TokenService;

    private final TokenProperties tokenProperties;

    private final ServerSecurityContextRepository securityContextRepository = NoOpServerSecurityContextRepository
            .getInstance();

    private final ReactiveAuthenticationSuccessHandler reactiveAuthenticationSuccessHandler = new ReactiveAuthenticationSuccessHandler() {
    };

    private final ReactiveAuthenticationFailureHandler reactiveAuthenticationFailureHandler = new ReactiveAuthenticationFailureHandler() {
    };


    @NonNull
    @SuppressWarnings("MVCPathVariableInspection")
    @PostMapping("${panda.spring.auth.endpoint.oauth-token:" + DEFAULT_PATH + "}")
    public Mono<Object> token(@PathVariable(REGISTRATION_ID) String registrationId,
                              @RequestBody OAuth2Request oAuth2Request,
                              ServerWebExchange serverWebExchange
    ) {
        return initAuthenticationToken(registrationId, oAuth2Request)
                // 使用认证链认证
                .flatMap(collaborativeReactiveAuthenticationStrategyResolver::resolve)
                // 保存上下文
                .flatMap(authenticationToken -> contextWrite(serverWebExchange, authenticationToken))
                // 生成token
                .flatMap(reactiveOAuth2TokenService::generate)
                // 成功的话，返回结果
                .flatMap(reactiveAuthenticationSuccessHandler::onAuthenticationSuccess)
                // 如果报错的话，返回错误信息
                .onErrorResume(reactiveAuthenticationFailureHandler::onAuthenticationFailure);
    }


    private Mono<CollaborativeAuthenticationToken> initAuthenticationToken(String registrationId, OAuth2Request oAuth2Request) {
        return reactiveClientRegistrationRepository.findByRegistrationId(registrationId)
                .map(clientRegistration -> {
                    if (StringKit.isBlank(oAuth2Request.getGrantType())) {
                        oAuth2Request.setGrantType(clientRegistration.getAuthorizationGrantType().getValue());
                    }
                    return new CollaborativeAuthenticationToken(oAuth2Request, clientRegistration);
                })
                .switchIfEmpty(Mono.error(new IllegalArgumentException("No clientRegistration found for registrationId: " + registrationId)));
    }

    /**
     * 保存上下文，先判断是否使用JWT，如果是JWT的话，就不需要保存了
     */
    private <U extends UserDetails> Mono<UserDetailsAuthenticationToken<U>> contextWrite(ServerWebExchange serverWebExchange, UserDetailsAuthenticationToken<U> authenticationToken) {
        // 判断是否是JWT，如果是的话就不保存了
        if (tokenProperties.getAccess().getMode() == TokenMode.JWT) {
            return Mono.just(authenticationToken);
        }
        SecurityContextImpl securityContext = new SecurityContextImpl(authenticationToken);
        return securityContextRepository.save(serverWebExchange, securityContext)
                .contextWrite(ReactiveSecurityContextHolder.withSecurityContext(Mono.just(securityContext)))
                .then(Mono.just(authenticationToken));
    }

}
